# pcap_monitor.py

import os
import time
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler

# 配置路径
PCAP_DIR = "/var/log/suricata/"  # 监控的 PCAP 文件目录
OUTPUT_DIR = "/var/log/suricata/"  # 流量特征的输出目录
FILE_STABILITY_TIME = 60  # 文件大小稳定时间，单位秒

# 用于存储新的 pcap 文件路径
pcapfiles = []

class PCAPHandler(FileSystemEventHandler):
    def on_created(self, event):
        if is_pcap_file(event.src_path):
            print(f"[INFO] New PCAP file detected: {event.src_path}")
            wait_for_file_stability(event.src_path)

            # 添加文件到 pcapfiles 列表
            pcapfiles.append(event.src_path)


def is_pcap_file(file_path):
    """检查文件是否为有效的 PCAP 文件"""
    if '.pcap' in os.path.basename(file_path):  # 检查文件名中是否包含 '.pcap'
        return True
    return False


def wait_for_file_stability(file_path):
    """等待文件大小在指定时间内不再变化"""
    last_size = None
    last_changed_time = time.time()

    while True:
        if os.path.exists(file_path):  # 检查文件是否存在
            current_size = os.path.getsize(file_path)
            current_time = time.time()

            # 如果文件大小不变且已经稳定了指定时间
            if current_size == last_size and (current_time - last_changed_time) >= FILE_STABILITY_TIME:
                print(f"[INFO] File size has stabilized. Processing...")
                break

            # 如果文件大小发生变化，更新最后修改时间和文件大小
            elif current_size != last_size:
                last_size = current_size
                last_changed_time = current_time
        else:
            print(f"[WARNING] File {file_path} no longer exists. Retrying...")

        # 等待 1 秒再次检查
        time.sleep(1)


def start_monitoring():
    """启动监控并返回监控到的 pcap 文件路径"""
    # 初始化 Watchdog 观察器
    event_handler = PCAPHandler()
    observer = Observer()
    observer.schedule(event_handler, PCAP_DIR, recursive=False)
    observer.start()

    print("[INFO] Watching PCAP directory...")

    try:
        while True:
            time.sleep(1)
            # 检查并返回 pcapfiles
            if pcapfiles:
                return pcapfiles  # 返回已经收集到的 PCAP 文件路径
    except KeyboardInterrupt:
        observer.stop()
    observer.join()
